import mplfinance as mpf
import pandas as pd
import numpy as np

# 获取数据，处理数据
def read_data(filename):
  df = pd.read_csv('./files/history_A_stock_k_data.csv')
  df['Open'] = df['open']
  df['Close'] = df['close']
  df['High'] = df['high']
  df['Low'] = df['low']
  df['Volume'] = df['volume']
  df['Date'] = df['date']

  df['Date'] = pd.to_datetime(df['Date'])
  df.set_index(['Date'], inplace=True)
  return df


# 获取数据
df = read_data('./files/test02.csv')

# 设置线元素的颜色
my_color = mpf.make_marketcolors(
  up="red",  # 上涨K线的颜色
  down="green",  # 下跌K线的颜色
  edge="black",  # 蜡烛图箱体的颜色
  volume="purple",  # 成交量柱子的颜色
  wick="black"  # 蜡烛图影线的颜色
)

# 自定义风格
my_style = mpf.make_mpf_style(
  base_mpf_style='blueskies',
  # base_mpl_style='seaborn',  # 也可以试试matplotlib的seaborn等风格。
  marketcolors=my_color,
  figcolor='(0, 0.8, 0.85)',
  gridcolor='(0.9, 0.9, 0.9)',
  rc={'font.family': 'SimHei', 'axes.unicode_minus': 'False'}
)

print(df.columns)
# 选择2023年7月11日到2024年5月17日的数据进行绘图
# mpf.plot(
  # df,
  # type='candle',
  # ylabel="open",
  # style=my_style,
  # title='sh600508 日线行情',
  # mav=(5, 10),
  # volume=True,
  # figratio=(5, 3),
  # ylabel_lower="vol", 
#   # savefig='sh600000.jpg' #保存为图片
# )





# data是测试数据，可以直接下载后读取，在下例中只显示其中100个交易日的数据
plot_data = df
# 读取显示区间最后一个交易日的数据
last_data = plot_data.iloc[-1]
# 使用mpf.figure()函数可以返回一个figure对象，从而进入External Axes Mode，从而实现对Axes对象和figure对象的自由控制
fig = mpf.figure(style=my_style, figsize=(12, 8), facecolor=(0.82, 0.83, 0.85))
# 添加三个图表，四个数字分别代表图表左下角在figure中的坐标，以及图表的宽（0.88）、高（0.60）
ax1 = fig.add_axes([0.06, 0.25, 0.88, 0.60])
# 添加第二、三张图表时，使用sharex关键字指明与ax1在x轴上对齐，且共用x轴
ax2 = fig.add_axes([0.06, 0.15, 0.88, 0.10], sharex=ax1)
ax3 = fig.add_axes([0.06, 0.05, 0.88, 0.10], sharex=ax1)
# 设置三张图表的Y轴标签
ax1.set_ylabel('price')
ax2.set_ylabel('volume')
ax3.set_ylabel('macd')
# 在figure对象上添加文本对象，用于显示各种价格和标题


fig.text(0.50, 0.94, 'sh.600508:--上海能源')

fig.text(0.12, 0.90, '开/收: ')
fig.text(0.17, 0.90, f'{np.round(last_data["open"], 3)} / {np.round(last_data["close"], 3)}')
fig.text(0.12, 0.86, '最后交易日期:')
fig.text(0.2, 0.86, f'{last_data.name.date()}')

fig.text(0.28, 0.90, '涨跌幅:')
# fig.text(0.33, 0.90, f'{last_data["change"]}')

fig.text(0.28, 0.86, '涨跌幅百分比:')
fig.text(0.36, 0.86, f'{np.round(last_data["pctChg"], 2)}%')


fig.text(0.40, 0.90, '高: ')
fig.text(0.42, 0.90, f'{last_data["high"]}')
fig.text(0.40, 0.86, '低: ')
fig.text(0.42, 0.86, f'{last_data["low"]}')
fig.text(0.55, 0.90, '量(股): ')
fig.text(0.60, 0.90, f'{last_data["volume"]}')
fig.text(0.55, 0.86, '额(元): ')
fig.text(0.60, 0.86, f'{last_data["amount"]}')
fig.text(0.70, 0.90, '涨停: ')
# fig.text(0.74, 0.90, f'{last_data["upper_lim"]}')
fig.text(0.70, 0.86, '跌停: ')
# fig.text(0.74, 0.86, f'{last_data["lower_lim"]}')
fig.text(0.85, 0.90, '均价: ')
# fig.text(0.89, 0.90, f'{np.round(last_data["average"], 3)}')
fig.text(0.85, 0.86, '昨收: ')
fig.text(0.89, 0.86, f'{last_data["preclose"]}')
# 调用mpf.plot()函数，注意调用的方式跟上一节不同，这里需要指定ax=ax1，volume=ax2，将K线图显示在ax1中，交易量显示在ax2中
mpf.plot(
  plot_data,
  ax=ax1,
  volume=ax2,
  type='candle',
  style=my_style,
  ylabel="价格",
  ylabel_lower="成交量", 
  figratio=(5, 3), # figratio设置图形的纵横比，figscale设置图像的缩放比例，tight_layout设置图像为紧密布局；

  # 使用mav关键字绘制移动平均线：
  # ①对单个移动平均线使用标量
  # ②对多个移动平均线使用元组或整数列表
  mav=(5, 10, 30),  
)
fig.show()		
input("Press Enter to close...")

